import os

from PIL import Image, ImageDraw


def generate_circle_pattern(image_size, circle_size, spacing, save_path):  # 生成mask图片
    image = Image.new('L', image_size, 0)
    draw = ImageDraw.Draw(image)

    width, height = image_size
    radius = circle_size // 2
    x_spacing = circle_size + spacing
    y_spacing = circle_size + spacing

    x_start = (width - (x_spacing * ((width - radius) // x_spacing + 1) - spacing)) // 2 + radius * 3
    y_start = (height - (y_spacing * ((height - radius) // y_spacing + 1) - spacing)) // 2 + radius * 2

    x = x_start
    while x < width:
        y = y_start
        while y < height:
            circle_position = (x - radius, y - radius, x + radius, y + radius)
            draw.ellipse(circle_position, fill=255)
            y += y_spacing
        x += x_spacing

    image.save(save_path)
    return image


def generate_mask_img():
    image_size = (320, 240)
    circle_size = 15
    spacing = 3
    save_path = r'./data/20230827/mask/mask.jpg'
    pattern_image = generate_circle_pattern(image_size, circle_size, spacing, save_path)
    pattern_image.show()  # 显示生成的图像


def apply_mask(input_image_path, mask_image):
    # 打开输入图像和Mask图像
    input_image = input_image_path.convert('L')

    # 确保Mask图像和输入图像具有相同的尺寸
    mask_image = mask_image.resize(input_image.size, Image.ANTIALIAS)

    # 创建一个与输入图像相同大小的透明图像
    result_image = Image.new('L', input_image.size)

    # 遍历每个像素，根据Mask图像中的像素值设置结果图像的像素值
    for x in range(input_image.width):
        for y in range(input_image.height):
            mask_pixel = mask_image.getpixel((x, y))
            if mask_pixel > 200:  # 白色
                result_pixel = input_image.getpixel((x, y))
            else:  # 黑色
                result_pixel = 0
            result_image.putpixel((x, y), result_pixel)
    return result_image


def process_train_imgs():  # 使用mask图片根据无遮挡图片生成遮挡图片
    ori_path = r'./data/20230827/ori'
    shade_path = r'./data/20230827/shade'
    mask_path = r'./data/20230827/mask/mask.jpg'
    mask_image = Image.open(mask_path)
    for file_name in os.listdir(ori_path):
        print(file_name)
        img_ori = Image.open(os.path.join(ori_path, file_name))
        img_shade = apply_mask(img_ori, mask_image)
        img_shade.save(os.path.join(shade_path, file_name))
    print('Done.')


if __name__ == '__main__':
    # generate_mask_img()
    process_train_imgs()
